#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep
import datetime

import grpc

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1
SWITCH_TO_SWITCH_PORT = 2

# 定义隧道ID映射
TUNNEL_IDS = {
    's1_s2_ingress': 100,  # s1->s2方向的入口隧道
    's2_s1_ingress': 101,  # s2->s1方向的入口隧道
    's1_s3_ingress': 200,  # s1->s3方向的入口隧道  
    's3_s1_ingress': 201,  # s3->s1方向的入口隧道
    's2_s3_ingress': 300,  # s2->s3方向的入口隧道
    's3_s2_ingress': 301,  # s3->s2方向的入口隧道
}

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr, t_port):
    """
    Installs three rules:
    1) An tunnel ingress rule on the ingress switch in the ipv4_lpm table that
       encapsulates traffic into a tunnel with the specified ID
    2) A transit rule on the ingress switch that forwards traffic based on
       the specified ID
    3) An tunnel egress rule on the egress switch that decapsulates traffic
       with the specified ID and sends it to the host

    :param p4info_helper: the P4Info helper
    :param ingress_sw: the ingress switch connection
    :param egress_sw: the egress switch connection
    :param tunnel_id: the specified tunnel ID
    :param dst_eth_addr: the destination IP to match in the ingress rule
    :param dst_ip_addr: the destination Ethernet address to write in the
                        egress rule
    :param is_direct_connection: whether switches are directly connected
    """
    # 1) Tunnel Ingress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)
    print("Installed ingress tunnel rule on %s for tunnel %d" % (ingress_sw.name, tunnel_id))

    # 2) Tunnel Transit Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": t_port
        })
    ingress_sw.WriteTableEntry(table_entry)
    print("Installed transit tunnel rule on %s for tunnel %d" % (ingress_sw.name, tunnel_id))

    # 3) Tunnel Egress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print("Installed egress tunnel rule on %s for tunnel %d" % (egress_sw.name, tunnel_id))

def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def readCounter(p4info_helper, sw, counter_name, index):
    """
    Reads the specified counter at the specified index from the switch.

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    :param counter_name: the name of the counter from the P4 program
    :param index: the counter index (tunnel ID)
    :return: counter data (packet_count, byte_count)
    """
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            return counter.data.packet_count, counter.data.byte_count
    return 0, 0

def writeCounterToFile(filename, switch_name, counter_name, tunnel_id, packet_count, byte_count):
    """
    Write counter information to file with timestamp
    """
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(filename, 'a') as f:
        f.write(f"[{timestamp}] {switch_name} {counter_name} {tunnel_id}: {packet_count} packets ({byte_count} bytes)\n")

def collectAndWriteCounters(p4info_helper, s1, s2, s3):
    """
    Collect all tunnel counters and write to respective files
    """
    print('\n----- Reading tunnel counters -----')
    
    # s1 -> s2 direction
    s1_ingress_packets, s1_ingress_bytes = readCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", TUNNEL_IDS['s1_s2_ingress'])
    s2_egress_packets, s2_egress_bytes = readCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", TUNNEL_IDS['s1_s2_ingress'])
    
    print("--- s1 -> s2 ---")
    print(f"s1 MyIngress.ingressTunnelCounter {TUNNEL_IDS['s1_s2_ingress']}: {s1_ingress_packets} packets ({s1_ingress_bytes} bytes)")
    print(f"s2 MyIngress.egressTunnelCounter {TUNNEL_IDS['s1_s2_ingress']}: {s2_egress_packets} packets ({s2_egress_bytes} bytes)")
    
    writeCounterToFile('S1S2.txt', 's1', 'MyIngress.ingressTunnelCounter', TUNNEL_IDS['s1_s2_ingress'], s1_ingress_packets, s1_ingress_bytes)
    writeCounterToFile('S1S2.txt', 's2', 'MyIngress.egressTunnelCounter', TUNNEL_IDS['s1_s2_ingress'], s2_egress_packets, s2_egress_bytes)
    
    # s2 -> s1 direction  
    s2_ingress_packets, s2_ingress_bytes = readCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", TUNNEL_IDS['s2_s1_ingress'])
    s1_egress_packets, s1_egress_bytes = readCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", TUNNEL_IDS['s2_s1_ingress'])
    
    print("--- s2 -> s1 ---")
    print(f"s2 MyIngress.ingressTunnelCounter {TUNNEL_IDS['s2_s1_ingress']}: {s2_ingress_packets} packets ({s2_ingress_bytes} bytes)")
    print(f"s1 MyIngress.egressTunnelCounter {TUNNEL_IDS['s2_s1_ingress']}: {s1_egress_packets} packets ({s1_egress_bytes} bytes)")
    
    writeCounterToFile('S1S2.txt', 's2', 'MyIngress.ingressTunnelCounter', TUNNEL_IDS['s2_s1_ingress'], s2_ingress_packets, s2_ingress_bytes)
    writeCounterToFile('S1S2.txt', 's1', 'MyIngress.egressTunnelCounter', TUNNEL_IDS['s2_s1_ingress'], s1_egress_packets, s1_egress_bytes)
    
    # s1 -> s3 direction
    s1_ingress_packets, s1_ingress_bytes = readCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", TUNNEL_IDS['s1_s3_ingress'])
    s3_egress_packets, s3_egress_bytes = readCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", TUNNEL_IDS['s1_s3_ingress'])
    
    print("--- s1 -> s3 ---")
    print(f"s1 MyIngress.ingressTunnelCounter {TUNNEL_IDS['s1_s3_ingress']}: {s1_ingress_packets} packets ({s1_ingress_bytes} bytes)")
    print(f"s3 MyIngress.egressTunnelCounter {TUNNEL_IDS['s1_s3_ingress']}: {s3_egress_packets} packets ({s3_egress_bytes} bytes)")
    
    writeCounterToFile('S1S3.txt', 's1', 'MyIngress.ingressTunnelCounter', TUNNEL_IDS['s1_s3_ingress'], s1_ingress_packets, s1_ingress_bytes)
    writeCounterToFile('S1S3.txt', 's3', 'MyIngress.egressTunnelCounter', TUNNEL_IDS['s1_s3_ingress'], s3_egress_packets, s3_egress_bytes)
    
    # s3 -> s1 direction
    s3_ingress_packets, s3_ingress_bytes = readCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", TUNNEL_IDS['s3_s1_ingress'])
    s1_egress_packets, s1_egress_bytes = readCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", TUNNEL_IDS['s3_s1_ingress'])
    
    print("--- s3 -> s1 ---")
    print(f"s3 MyIngress.ingressTunnelCounter {TUNNEL_IDS['s3_s1_ingress']}: {s3_ingress_packets} packets ({s3_ingress_bytes} bytes)")
    print(f"s1 MyIngress.egressTunnelCounter {TUNNEL_IDS['s3_s1_ingress']}: {s1_egress_packets} packets ({s1_egress_bytes} bytes)")
    
    writeCounterToFile('S1S3.txt', 's3', 'MyIngress.ingressTunnelCounter', TUNNEL_IDS['s3_s1_ingress'], s3_ingress_packets, s3_ingress_bytes)
    writeCounterToFile('S1S3.txt', 's1', 'MyIngress.egressTunnelCounter', TUNNEL_IDS['s3_s1_ingress'], s1_egress_packets, s1_egress_bytes)
    
    # s2 -> s3 direction
    s2_ingress_packets, s2_ingress_bytes = readCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", TUNNEL_IDS['s2_s3_ingress'])
    s3_egress_packets, s3_egress_bytes = readCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", TUNNEL_IDS['s2_s3_ingress'])
    
    print("--- s2 -> s3 ---")
    print(f"s2 MyIngress.ingressTunnelCounter {TUNNEL_IDS['s2_s3_ingress']}: {s2_ingress_packets} packets ({s2_ingress_bytes} bytes)")
    print(f"s3 MyIngress.egressTunnelCounter {TUNNEL_IDS['s2_s3_ingress']}: {s3_egress_packets} packets ({s3_egress_bytes} bytes)")
    
    writeCounterToFile('S2S3.txt', 's2', 'MyIngress.ingressTunnelCounter', TUNNEL_IDS['s2_s3_ingress'], s2_ingress_packets, s2_ingress_bytes)
    writeCounterToFile('S2S3.txt', 's3', 'MyIngress.egressTunnelCounter', TUNNEL_IDS['s2_s3_ingress'], s3_egress_packets, s3_egress_bytes)
    
    # s3 -> s2 direction
    s3_ingress_packets, s3_ingress_bytes = readCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", TUNNEL_IDS['s3_s2_ingress'])
    s2_egress_packets, s2_egress_bytes = readCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", TUNNEL_IDS['s3_s2_ingress'])
    
    print("--- s3 -> s2 ---")
    print(f"s3 MyIngress.ingressTunnelCounter {TUNNEL_IDS['s3_s2_ingress']}: {s3_ingress_packets} packets ({s3_ingress_bytes} bytes)")
    print(f"s2 MyIngress.egressTunnelCounter {TUNNEL_IDS['s3_s2_ingress']}: {s2_egress_packets} packets ({s2_egress_bytes} bytes)")
    
    writeCounterToFile('S2S3.txt', 's3', 'MyIngress.ingressTunnelCounter', TUNNEL_IDS['s3_s2_ingress'], s3_ingress_packets, s3_ingress_bytes)
    writeCounterToFile('S2S3.txt', 's2', 'MyIngress.egressTunnelCounter', TUNNEL_IDS['s3_s2_ingress'], s2_egress_packets, s2_egress_bytes)
    
    print("--- Finished ---")

def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))

def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    # Clear existing counter files
    for filename in ['S1S2.txt', 'S1S3.txt', 'S2S3.txt']:
        if os.path.exists(filename):
            os.remove(filename)

    try:
        # Create switch connections for s1, s2, and s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', 
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on all switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        # Write tunnel rules for all host pairs
        
        # h1 (10.0.1.1) <-> h2 (10.0.2.2) via s1-s2
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=TUNNEL_IDS['s1_s2_ingress'],
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2", t_port=2)
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=TUNNEL_IDS['s2_s1_ingress'],
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1", t_port=2)

        # h1 (10.0.1.1) <-> h3 (10.0.3.3) via s1-s3  
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=TUNNEL_IDS['s1_s3_ingress'],
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3", t_port=3)
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=TUNNEL_IDS['s3_s1_ingress'],
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1", t_port=2)

        # h2 (10.0.2.2) <-> h3 (10.0.3.3) via s2-s3
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=TUNNEL_IDS['s2_s3_ingress'],
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3", t_port=3)
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=TUNNEL_IDS['s3_s2_ingress'],
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2", t_port=3)

        # Read table rules (optional)
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2) 
        readTableRules(p4info_helper, s3)

        # Continuously collect and write counters
        while True:
            sleep(2)
            collectAndWriteCounters(p4info_helper, s1, s2, s3)

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)